# import cv2
from cgi import test
from importlib_metadata import distribution, re
import numpy as np
import math
import time
import sys
import os
import argparse
from colorama import Fore, Back, Style
from colorama import init

###############################################################################
# 1D wave function collapse
###############################################################################

table_capacity = 20
collapsed = []
entropy = []
candidates = [1,2,3,4,5]
candidates_weight = [1.0,2.0,3.0,4.0,2.0]

# the wave
# wave is a 2D array of size (table_capacity, number_of_candidates)
wave = []

# all the compatible situations
rules = {
    (1,2),
    (2,3),
    (3,4),
    (4,5),
    (5,4),
    (4,3),
    (3,2),
    (2,1),
    (1,5),
    (5,1),
    }

# contradiction check
def contradiction_check(wave):

    for i in range(len(wave)):

        neibhor_1 = i-1
        neibhor_2 = i+1

        if i == len(wave)-1:
            neibhor_2 = 0
        if i == 0:
            neibhor_1 = len(wave)-1

        for element_i in wave[i]:
            for element_j in wave[neibhor_2]:
                if not (element_i, element_j) in rules:
                    return False

        for element_i in wave[i]:
            for element_j in wave[neibhor_1]:
                if not (element_j, element_i) in rules:
                    return False

    return True

# distributional trigger
def distributional_trigger(wavelet_index):

    collapse_element_index = 0

    # construct a new weight array for the wavelet
    weight_array = []
    for i in range(len(wave[wavelet_index])):
        weight_array.append(get_weight_of_candidate(wave[wavelet_index][i]))

    # generate a random number in [0, sum(weight_array)]
    random_number = np.random.rand() * sum(weight_array)

    # find the collapse element
    for i in range(len(weight_array)):
        random_number -= weight_array[i]
        if random_number <= 0:
            collapse_element_index = i
            break

    return collapse_element_index

# get weight through candidate
def get_weight_of_candidate(candidate):
    for i in range(len(candidates)):
        if candidates[i] == candidate:
            return candidates_weight[i]

# calculate the Shannon entropy of a wavelet
def shannon_entropy_of_wavelet(wavelet):

    entropy = 0
    sum_weight = 0
    for i in wavelet:
        sum_weight += get_weight_of_candidate(i)

    for i in wavelet:
        probability = get_weight_of_candidate(i) / sum_weight
        probability = max(probability, 1e-10)
        entropy += -probability * math.log(probability)

    # add a small purturbation
    entropy += np.random.rand() * 1e-3

    return entropy

# find an element to collapse
def find_minimum_entropy_position():
    # find minimum entropy index of entropy
    min_entropy_index = 0
    min_entropy = 1e10
    for i in range(len(entropy)):
        if entropy[i] < min_entropy and collapsed[i] == False:
            min_entropy = entropy[i]
            min_entropy_index = i

    return min_entropy_index

# collapse the wavelet
def collapse_at_position(wavelet_index):
    
    collapse_element_index = distributional_trigger(wavelet_index)
    wavelet = []
    wavelet.append(wave[wavelet_index][collapse_element_index])
    wave[wavelet_index] = wavelet.copy()
    if len(wave[wavelet_index]) == 1:
        collapsed[wavelet_index] = True


# update entropy
def update_wave_entropy():
    for i in range(len(wave)):
        entropy[i] = shannon_entropy_of_wavelet(wave[i])

# propagate the wave function
def propagate(wave):

    for k in range(5):
        for i in range(len(wave)):

            if len(wave[i]) == len(candidates):
                continue

            neibhor_1 = i-1
            neibhor_2 = i+1

            if i == len(wave)-1:
                neibhor_2 = 0
            if i == 0:
                neibhor_1 = len(wave)-1

            # constrain neibhors with the rules
            wavelet_temp1 = set()
            wavelet_temp2 = set()
            for element_i in wave[i]:
                for element_j in wave[neibhor_2]:
                    if len(wave[neibhor_2]) == 1:
                        collapsed[neibhor_2] = True
                        break
                    if (element_i, element_j) in rules:
                        wavelet_temp2.add(element_j)

            for element_i in wave[i]:
                for element_j in wave[neibhor_1]:
                    if len(wave[neibhor_1]) == 1:
                        collapsed[neibhor_1] = True
                        break
                    if (element_j, element_i) in rules:
                        wavelet_temp1.add(element_j)

            if len(wavelet_temp2) != 0:
                wave[neibhor_2].clear()
                for element_i in wavelet_temp2:
                    wave[neibhor_2].append(element_i)
            
            if len(wavelet_temp1) != 0:
                wave[neibhor_1].clear()
                for element_i in wavelet_temp1:
                    wave[neibhor_1].append(element_i)
    
    update_wave_entropy()

# check if the wave is collapsed
def is_collapsed(wave):
    for i in range(len(wave)):
        if collapsed[i] == False:
            return False
    return True

# main function
def main():

    # input('Press Enter to continue...')

    print(Fore.GREEN + "wfc_toy" + Style.RESET_ALL)

    for i in range(table_capacity):
        entropy.append(1)
        collapsed.append(False)

    # initialieze the wave
    for i in range(table_capacity):
        # make a copy of candidates
        wave.append(candidates.copy())

    # seed the random number generator
    np.random.seed(int(time.time()))

    update_wave_entropy()

    # collapse the wave
    while not is_collapsed(wave):

        # 1. select the minimum entropy wavelet 
        collapse_element = find_minimum_entropy_position()

        # 2. collapse the minimum entropy wavelet
        collapse_at_position(collapse_element)
        
        # 3. propagate the wave function
        propagate(wave)

        print('\033c',end='')
        for i in range(len(wave)):
            # convert the wavelet to string
            wavelet_str = str(wave[i])
            # convert entropy to string
            entropy_str = str(entropy[i])
            if collapsed[i] == True:
                print(Fore.RED + wavelet_str + " , " + entropy_str + Style.RESET_ALL)
            else:
                print(wavelet_str + " , " + entropy_str)
        print("-----------------------------------------------------")

        # keyboard input
        input()

        # delay 100ms
        time.sleep(0.1)

    # compatability check
    if contradiction_check(wave) == True:
        print('\033c',end='')
        for i in range(len(wave)):
            # convert the wavelet to string
            wavelet_str = str(wave[i])
            if collapsed[i] == True:
                print(Fore.GREEN + wavelet_str + Style.RESET_ALL)
            else:
                print(wavelet_str + " , " + entropy_str)
        print("-----------------------------------------------------")
        print("No contradiction found!")
    else:
        print("contradiction !")

# main entry
if __name__ == "__main__":
    main()
